import cv2
import sys
import numpy as np
import datetime
import os
import glob
from retinaface import RetinaFace
import time


def retinaface_detection(img):
	""" 基于RetinaFace模型进行人脸检测
		img: 待人脸检测图像
		return: faces and landmarks
	"""
	start_time = time.clock()

	# 定义相关参数
	thresh = 0.8 
	scales = [1024, 1980]
	count = 1
	gpuid = 0 # 设置为-1为使用CPU

	# 加载RetinaFace模型
	detector = RetinaFace('./model/retinaface-R50-cudnnoff/R50', 0, gpuid, 'net3')

	# count time of loading model
	modelload_time = time.clock()
	print('Load model time: %s s' % (modelload_time - start_time))
	
	#print(img.shape)
	im_shape = img.shape
	target_size = scales[0]
	max_size = scales[1]
	im_size_min = np.min(im_shape[0:2])
	im_size_max = np.max(im_shape[0:2])

	#im_scale = 1.0
	#if im_size_min>target_size or im_size_max>max_size:
	im_scale = float(target_size) / float(im_size_min)
	# prevent bigger axis from being more than max_size:
	if np.round(im_scale * im_size_max) > max_size:
		im_scale = float(max_size) / float(im_size_max)

	print('im_scale', im_scale)
	scales = [im_scale]
	flip = False

	setting_time = time.clock()
	print('Setting parameters time: %s s' % (setting_time - modelload_time))

	for c in range(count):
		faces, landmarks = detector.detect(img,thresh,scales=scales,do_flip=flip)
		print(c, faces.shape, landmarks.shape)

	facedetection_time = time.clock()
	print('Face detection time: %s s' % (facedetection_time - setting_time))

	return faces, landmarks
	"""
	if faces is not None:
		print('find', faces.shape[0], 'faces')
		for i in range(faces.shape[0]):
			#print('score', faces[i][4])
			box = faces[i].astype(np.int)
			#color = (255,0,0)
			color = (0, 0, 255)
			cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), color, 2)
			if landmarks is not None:
				landmark5 = landmarks[i].astype(np.int)
            	#print(landmark.shape)
				for l in range(landmark5.shape[0]):
					color = (0, 0, 255)
					if l == 0 or l == 3:
						color = (0, 255, 0)
					cv2.circle(img, (landmark5[l][0], landmark5[l][1]), 1, color, 2)
        
		filename = './detector_test.jpg'
		print('writing', filename)
		cv2.imwrite(filename, img)
		"""


if __name__ == '__main__':
	
	img = cv2.imread('imgs/t1.jpg')
	faces, landmarks = retinaface_detection(img)

	start_time0 = time.clock()

	if faces is not None:
		print('find', faces.shape[0], 'faces')
		for i in range(faces.shape[0]):
			#print('score', faces[i][4])
			box = faces[i].astype(np.int)
			#color = (255,0,0)
			color = (0, 0, 255)
			cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), color, 2)
			if landmarks is not None:
				landmark5 = landmarks[i].astype(np.int)
            	#print(landmark.shape)
				for l in range(landmark5.shape[0]):
					color = (0, 0, 255)
					if l == 0 or l == 3:
						color = (0, 255, 0)
					cv2.circle(img, (landmark5[l][0], landmark5[l][1]), 1, color, 2)

		draw_time = time.clock()
		print('Draw circle time: %s s' % (draw_time - start_time0))
        
		filename = './ret/detector_test.jpg'
		print('writing', filename)
		cv2.imwrite(filename, img)

		saveimg_time = time.clock()
		print('Save img time: %s s' % (saveimg_time - draw_time))













